package pac;

import io.PlinkFileException;
import io.PlinkFilereader;

import java.io.IOException;
import java.text.DecimalFormat;
import java.util.BitSet;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Vector;

import preprocessing.WindowingSystem;
import statistics.SignificanceHelper;
import datastructures.CarrierIndex;
import datastructures.Individual;
import datastructures.IndividualInfo;
import datastructures.SNPInfo;
import datastructures.UserParameters;


public class SearchDatastructure {

	//publicly accessible
	public final long seed = 1000;
	public long startTimeMs;
 	public volatile long stage1Candidates;
 	public volatile long stage1Trials;
	public volatile long stage1Time;
	public volatile long stage2Time;
	
	public final UserParameters userParams;
	public WindowingSystem windowMaker;
	
	//publicly accessible MAP file info
	public SNPInfo[] snpList;
	public BitSet[] snpMasks;
	public CarrierIndex snpIndices;
	
	//publicly accessible RAW/PED file info
	public IndividualInfo[] caseInfo;
	public IndividualInfo[] controlInfo;
	public BitSet[] caseGenotypes; 
	public BitSet[] controlGenotypes;

	// singleton shared accross threads
	public final SignificanceHelper significance; 
	
	//publicly accessible datastructure visible to the threads
	public volatile Map<String, Integer> shortlistedCandidates = Collections.synchronizedMap (new HashMap<String, Integer>());
	
	/*
	 * Constructor - loads the datafiles
	 */
	public SearchDatastructure ( UserParameters userParameters, WindowingSystem windowingSystem, SignificanceHelper significanceHelper ) throws Exception {
		
		userParams = userParameters;
		windowMaker = windowingSystem;
		significance = significanceHelper;
		
		// ----- The sequence of the following operations is important ----- //
		
		//1. Read in the cases (for sampling) and controls (expected) data
		getCasesAndControls();
		
		//2. then read map 
		getMapFile();
		
		System.out.println(caseInfo.length + " cases, " + controlInfo.length + " controls.");
		DecimalFormat doubleF	= new DecimalFormat("#.###"); 
		System.out.println(snpList.length + " SNPs. Testing 4 interaction models per SNP-pair.");
		System.out.println("Approximately " + doubleF.format( (double) snpList.length*(snpList.length-1)*0.5*0.000001*4) + " million pairwise tests. \n");
		
		//3. Now set the number of cases and controls, and calculate significance cutoffs.
		System.out.println("Calculating significance vectors ... "); 
		significance.calculatePvalCutoffs(userParameters);
		
		//4. index SNPs by carrier status
		assignCarrierStatusToSNPs();
		
		//5. then create the SNP Masks that put snp-variables into bins.
		assignSNPsToMasks();
		
		//6. read LD block size information
		setupLDBlockSize();
		
		//Check the heap state, and quit.
		Runtime r = Runtime.getRuntime();
		r.gc();
		long used_mem = r.totalMemory() - r.freeMemory();
		System.out.println("Heap space used up by dataset: " + (float)used_mem/1000000 + " Mb");	
		
	}
	
	
	private void setupLDBlockSize() throws IOException, PlinkFileException {
		PlinkFilereader filereader = new PlinkFilereader();
		//get the block size we are looking across.
		if ( userParams.getBlockSizeBp() != -1)
			padSNPBlocks(userParams.getBlockSizeBp(), 1);
		else if ( userParams.getBlockSizeCm() != -1)
			padSNPBlocks(userParams.getBlockSizeCm(), 2);
		else {
			boolean foundBlockFile = filereader.readBlockFile( userParams.getBlockfilename(), snpList );
			// ? is this correct?
			if( ! foundBlockFile || (userParams.getBlockSizeBp() != -1 && userParams.getBlockSizeCm() != -1 )) {
				padSNPBlocks(0, 3);
			}
		}
	}
		
	
	private void assignSNPsToMasks() throws IOException, PlinkFileException 
	{
		windowMaker.makeWindowedGenomesFromData(this);
		
		//flush the memory
		caseGenotypes = null;
		controlGenotypes = null;
		System.gc();
	}
	
	
	/*
	 * Cross indexing allele carriers by SNP.
	 */
	private void assignCarrierStatusToSNPs() {
		System.out.println("Indexing datasets (this may take a while) ... ");
		
		int numberOfCases = caseGenotypes.length;
		int numberOfControls = controlGenotypes.length;
		
		int numberOfSNPs = snpList.length;
		snpIndices = new CarrierIndex(numberOfSNPs * 2); 
		
		for (int snp=0; snp<snpList.length; snp++ ) {			
			BitSet caseDomEntries = new BitSet(numberOfCases);
			BitSet caseRecEntries = new BitSet(numberOfCases);
			
			for(int indiv=0; indiv<numberOfCases; indiv++) {
				if( caseGenotypes[indiv].get(2*snp+1) ) {
					caseDomEntries.set(indiv); caseRecEntries.set(indiv);
				}
				else if ( caseGenotypes[indiv].get(2*snp) ) {
					caseDomEntries.set(indiv);
				}
			}
			snpIndices.setCaseCarrierIndex(caseDomEntries, 2*snp);
			snpIndices.setCaseCarrierIndex(caseRecEntries, 2*snp+1);
			
			BitSet controlDomEntries = new BitSet(numberOfControls);
			BitSet controlRecEntries = new BitSet(numberOfControls);
			for(int indiv=0; indiv<numberOfControls; indiv++) {
				if( controlGenotypes[indiv].get(2*snp+1) ) {
					controlDomEntries.set(indiv); controlRecEntries.set(indiv);
				}
				else if ( controlGenotypes[indiv].get(2*snp) ) {
					controlDomEntries.set(indiv);
				}
			}
			snpIndices.setControlCarrierIndex(controlDomEntries, 2*snp);
			snpIndices.setControlCarrierIndex(controlRecEntries, 2*snp+1);			
		}
		
	}
	
	
	private void getMapFile() throws IOException, PlinkFileException {
		PlinkFilereader filereader = new PlinkFilereader();
		snpList = filereader.readMapfile ( userParams.getMapfilename() );
		if ( snpList.length != caseInfo[0].getNumberOfSNPs() )
			throw new PlinkFileException("Raw and Map files are incompatible");
		userParams.setNumber_of_snps(snpList.length);
	}
	
	
	/*
	 * Populates caseInfo, controlInfo, caseGenotypes and controlGenotypes
	 * Deals with the nuances of reading raw files and segregating cases from controls.
	 */
	private void getCasesAndControls() throws IOException, PlinkFileException 
	{	
		PlinkFilereader filereader = new PlinkFilereader();
		Individual[] tempCases;
		Individual[] tempControls;
		
		//1. If a single raw file has been provided.
		if(userParams.getAllRawfilename().compareTo("")!=0) 
		{	
			Individual[] allIndividuals = filereader.readRawFile( userParams.getAllRawfilename() );
			Vector<Individual[]> casesAndControls = filereader.segregateCasesAndControls(allIndividuals, userParams.getPhenofilename(), false);
			tempCases = casesAndControls.elementAt(0);
			tempControls = casesAndControls.elementAt(1); 
		}
		
		//2. Else read separate files.
		else 
		{
			tempCases =  filereader.readRawFile( userParams.getCasefilename() );
			tempControls = filereader.readRawFile( userParams.getControlfilename() );
		}
		
		//3. retain only the required datastructures.
		caseInfo = new IndividualInfo [tempCases.length];
		controlInfo = new IndividualInfo [tempControls.length];
		
		caseGenotypes = new BitSet [tempCases.length];		
		for(int i=0; i<tempCases.length; i++) {
			caseInfo[i] = tempCases[i].getIndividualInfo();
			caseGenotypes[i] = tempCases[i].getGenotype();
		}
		userParams.setNumber_of_cases ( caseGenotypes.length );
		
		//optionally, if specified, read in the controls
		if (tempControls != null) {
			controlGenotypes = new BitSet [tempControls.length];	
			for(int i=0; i<tempControls.length; i++) {
				controlInfo[i] = tempControls[i].getIndividualInfo();
				controlGenotypes[i] = tempControls[i].getGenotype();
			}
			userParams.setNumber_of_controls( controlGenotypes.length );
		}
		else
			controlGenotypes = null;
		
		
	}
	
	/*
	 * Helper method: in absence of block file, either --useblocks or each SNP is its own block
	 */
	private void padSNPBlocks (float windowSize, int mode) {
		
		//block file provided
		if(mode == 3) {
			for(int i=0; i<snpList.length; i++) 
				snpList[i].setBlockNo(i);
		}
		
		//cm distance
		else if (mode == 2) {
			for(int i=0; i<snpList.length; i++) {
				double offset = (double) snpList[i].getCm_distance() / windowSize;
				double base = (double) snpList[i].getChromosome() * 100000; //currently using 100,000 cM
				double windowNo =  base + offset;
				snpList[i].setBlockNo(windowNo);
			}
		}
		
		//bp distance
		else if (mode ==1 ) {
			for(int i=0; i<snpList.length; i++) {
				float offset = (float) snpList[i].getBp_distance() / windowSize;
				float base = (float) snpList[i].getChromosome() * (200000000 / windowSize); //currently using 200M for chr1 size
				float windowNo =  base + offset;
				snpList[i].setBlockNo(windowNo);
			}
		}
		
	}	

}
